from .stereo_transformer import STEREOTrans
import numpy as np
import ctypes
import torch
# Initialize PyCUDA
class STEREOTransTRT(STEREOTrans):
    def __init__(self, model_cfg, num_class, dataset):
        super().__init__(model_cfg, num_class, dataset)
        self.module_topology.remove('dense_head_2d')
        self.module_topology.remove('depth_loss_head')
        self.module_topology.remove('voxel_loss_head')
        self.module_list = self.build_networks(trt=True)
        self.engine = None
        
    def load_engine(self, engine_path):
        import tensorrt as trt
        import pycuda.autoinit
        import pycuda.driver as cuda
        cuda.init()
        TRT_LOGGER = trt.Logger()
        device = cuda.Device(0)
        context = device.make_context()
        ctypes.CDLL("/cv/lyh/trt_models/tensorrt_dsgn2/plugins/libCostVolumePlugin.so", mode=ctypes.RTLD_GLOBAL)
        ctypes.CDLL("/cv/lyh/trt_models/tensorrt_dsgn2/plugins/libDeformConvPlugin.so", mode=ctypes.RTLD_GLOBAL)
        ctypes.CDLL("/cv/lyh/trt_models/tensorrt_dsgn2/build/libGridSamplePlugin.so", mode=ctypes.RTLD_GLOBAL)
        self.engine = TrtModel(engine_path)
        print('Stereo transformer engine loaded successfully')
        
    def forward(self, *inputs):
        if self.engine:
            # batch_dict = inputs[0]
            
            return self.forward_engine(*inputs)
        
        batch_dict = self.backbone_3d(*inputs)
        # left_stereo_features = batch_dict['left_stereo_feat']
        # right_stereo_features = batch_dict['right_stereo_feat']
        # sem_features = batch_dict['sem_features'][-3]
        # return left_stereo_features, right_stereo_features, sem_features
        for cur_module in self.module_list[1:]:
            batch_dict = cur_module(batch_dict)
            
        output_dict = self.dense_head.forward_ret_dict
        cls_preds = output_dict['cls_preds']
        box_preds = output_dict['box_preds']
        dir_cls_preds = output_dict['dir_cls_preds']
        # batch_cls_preds = batch_dict['batch_cls_preds']
        # batch_box_preds = batch_dict['batch_box_preds']
        return dir_cls_preds, cls_preds, box_preds
        # return dir_cls_preds, cls_preds, box_preds, batch_cls_preds, batch_box_preds
    
    def forward_engine(self, *inputs):
        box_preds, cls_preds, dir_cls_preds = self.engine(*inputs)
        box_preds = torch.from_numpy(box_preds.reshape(1,152,144,98)).cuda()
        cls_preds = torch.from_numpy(cls_preds.reshape(1,152,144,98)).cuda()
        dir_cls_preds = torch.from_numpy(dir_cls_preds.reshape(1,152,144,28)).cuda()
        data_dict = {}
        data_dict['batch_size'] = 1
        data_dict['cls_preds'] = cls_preds
        data_dict['box_preds'] = box_preds
        data_dict['dir_cls_preds'] = dir_cls_preds
        
        batch_cls_preds, batch_box_preds = self.dense_head.generate_predicted_boxes(batch_size=1, 
                                                 cls_preds=cls_preds, 
                                                 box_preds=box_preds, 
                                                 dir_cls_preds=dir_cls_preds)
        data_dict['batch_cls_preds'] = batch_cls_preds
        data_dict['batch_box_preds'] = batch_box_preds
        # TODO: check the code here, we add sigmoid in the generate predicted boxes, so set normalized to be True
        data_dict['cls_preds_normalized'] = False
        pred_dicts, ret_dicts = self.post_processing(data_dict)
        return pred_dicts, ret_dicts
        
        
        
class TrtModel:

    def __init__(self,engine_path,max_batch_size=1,dtype=np.float32):

        self.engine_path = engine_path
        self.dtype = dtype
        self.logger = trt.Logger(trt.Logger.WARNING) # it's ok. runtime import
        self.runtime = trt.Runtime(self.logger)
        self.engine = self.load_engine(self.runtime, self.engine_path)
        self.max_batch_size = max_batch_size
        self.inputs, self.outputs, self.bindings, self.stream = self.allocate_buffers()
        self.context = self.engine.create_execution_context()

    @staticmethod
    def load_engine(trt_runtime, engine_path):
        # trt.init_libnvinfer_plugins(None, "")         
        trt.init_libnvinfer_plugins(TRT_LOGGER, '')    
        with open(engine_path, 'rb') as f:
            engine_data = f.read()
        engine = trt_runtime.deserialize_cuda_engine(engine_data)
        return engine

    def allocate_buffers(self):

        inputs = []
        outputs = []
        bindings = []
        stream = cuda.Stream()

        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding)) * self.max_batch_size

            host_mem = cuda.pagelocked_empty(size, self.dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)

            bindings.append(int(device_mem))

            if self.engine.binding_is_input(binding):
                inputs.append(HostDeviceMem(host_mem, device_mem))
            else:
                outputs.append(HostDeviceMem(host_mem, device_mem))

        return inputs, outputs, bindings, stream


    def __call__(self, batched_img, downsampled_disp, norm_coord_imgs, valids, batch_size=1):

        batched_img = batched_img.astype(self.dtype)
        downsampled_disp = downsampled_disp.astype(self.dtype)
        norm_coord_imgs = norm_coord_imgs.astype(self.dtype)
        valids = valids.astype(self.dtype)

        np.copyto(self.inputs[0].host,batched_img.ravel())
        np.copyto(self.inputs[1].host,downsampled_disp.ravel())
        np.copyto(self.inputs[2].host, norm_coord_imgs.ravel())
        np.copyto(self.inputs[3].host, valids.ravel())

        for inp in self.inputs:
            cuda.memcpy_htod_async(inp.device, inp.host, self.stream)

        # self.context.execute_async(batch_size=batch_size, bindings=self.bindings, stream_handle=self.stream.handle)
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        for out in self.outputs:
            cuda.memcpy_dtoh_async(out.host, out.device, self.stream) 


        self.stream.synchronize()
        return [out.host for out in self.outputs]
    
class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()
